import tensorflow as tf
import numpy as np
import random
from collections import deque

# Hyper Parameters for DQN
GAMMA = 0.7  # discount factor for target Q
INITIAL_EPSILON = 0.7  # starting value of epsilon
FINAL_EPSILON = 0.01  # final value of epsilon
REPLAY_SIZE = 10000  # experience replay buffer size
BATCH_SIZE = 64  # size of minibatch


class DQN():
    def __init__(self, board_width=8,board_height=8,model_file=None):
        # init experience replay
        self.replay_buffer = deque()
        # init some parameters
        self.time_step = 0

        self.epsilon = INITIAL_EPSILON

        self.board_width=board_width
        self.board_height=board_height

        self.state_dim = self.board_width * self.board_height

        self.action_dim = self.board_width * self.board_height

        self.create_Q_network()
        self.create_training_method()

        # Init session
        self.session = tf.InteractiveSession()
        self.session.run(tf.initialize_all_variables())

        self.saver = tf.train.Saver()
        if model_file is not None:
            self.restore_model(model_file)

    def create_Q_network(self, board_width=8, board_height=8):

        self.input_states = tf.placeholder(
            tf.float32,
            shape=[None, 4,self.board_height, self.board_width])  # [batch, in_height, in_width, in_channels]这样的shape

        #tf.reshape(self.input_states, [-1, board_height, board_width, 4])
        self.input_state = tf.transpose(self.input_states, [0, 2, 3, 1])

        ##########################################################

        W_conv1 = self.weight_variable([3, 3, 4, 32])  # 前两个维度是patch的大小，接着是输入的通道数目，最后是输出的通道数目

        self.conv1 = self.conv2d(self.input_state, W_conv1)

        self.h_conv1 = tf.nn.relu(self.conv1)

        ##########################################################

        ##########################################################

        W_conv2 = self.weight_variable([3, 3, 32, 64])  # 前两个维度是patch的大小，接着是输入的通道数目，最后是输出的通道数目

        self.conv2 = self.conv2d(self.h_conv1, W_conv2)

        self.h_conv2 = tf.nn.relu(self.conv2)

        ##########################################################

        ##########################################################
        W_conv2_1 = self.weight_variable([3, 3, 64, 128])  # 前两个维度是patch的大小，接着是输入的通道数目，最后是输出的通道数目

        self.conv2_1 = self.conv2d(self.h_conv2, W_conv2_1)

        self.h_conv2_1 = tf.nn.relu(self.conv2_1)

        W_conv3 = self.weight_variable([3, 3, 128, 256])  # 前两个维度是patch的大小，接着是输入的通道数目，最后是输出的通道数目

        self.conv3 = self.conv2d(self.h_conv2_1, W_conv3)

        self.h_conv3 = tf.nn.relu(self.conv3)

        # W_conv3 = self.weight_variable([3, 3, 64, 128])  # 前两个维度是patch的大小，接着是输入的通道数目，最后是输出的通道数目
        #
        # self.conv3 = self.conv2d(self.h_conv2, W_conv3)
        #
        # self.h_conv3 = tf.nn.relu(self.conv3)

        ##########################################################

        ##########################################################

        # 3-1 Action Networks

        ##########################################################


        W_conv4 = self.weight_variable([1, 1, 256, 4])  # 前两个维度是patch的大小，接着是输入的通道数目，最后是输出的通道数目

        self.conv4 = self.conv2d(self.h_conv3, W_conv4)

        self.h_conv4 = tf.nn.relu(self.conv4)

        self.action_conv_flat = tf.reshape(
            self.h_conv4, [-1, 4 * self.board_height * self.board_width])


        self.action_fc = tf.layers.dense(inputs=self.action_conv_flat,  # 全连接
                                         units=self.board_height * self.board_width,
                                         activation=tf.nn.sigmoid)#tf.nn.log_softmax)


        self.Q_value = self.action_fc

        self.evaluation_conv_flat = tf.reshape(
            self.h_conv4, [-1, 4 * board_height * board_width])
        self.evaluation_fc1 = tf.layers.dense(inputs=self.evaluation_conv_flat,
                                              units=64, activation=tf.nn.relu)
        self.evaluation_fc2 = tf.layers.dense(inputs=self.evaluation_fc1,
                                              units=1, activation=tf.nn.tanh)

        ##########################################################

    def create_training_method(self):
        self.labels = tf.placeholder(tf.float32, shape=[None, 1])
        self.value_loss = tf.losses.mean_squared_error(self.labels,
                                                       self.evaluation_fc2) # 均方误差

        self.action_input = tf.placeholder("float", [None,self.action_dim])  # one hot presentation
        self.y_input = tf.placeholder("float", [None])

        Q_action = tf.reduce_sum(tf.multiply(self.Q_value, self.action_input), reduction_indices=1) # 行求和
        self.cost =tf.reduce_mean(tf.square(self.y_input - Q_action))+self.value_loss
        self.optimizer = tf.train.AdamOptimizer(0.00001).minimize(self.cost)

        # self.learning_rate = 0.001
        #
        # self.mcts_probs = tf.placeholder(
        #     tf.float32, shape=[None, self.board_height * self.board_width])
        #
        # self.policy_loss = tf.negative(tf.reduce_mean(
        #     tf.reduce_sum(tf.multiply(self.mcts_probs, self.action_fc), 1)))
        #
        # self.optimizer = tf.train.AdamOptimizer(
        #     learning_rate=self.learning_rate).minimize(self.policy_loss)


    def perceive(self, state, action, reward, next_state, done,new_game_idx,winner):
        one_hot_action = np.zeros(self.action_dim)
        one_hot_action[action] = 1
        self.replay_buffer.append([state, one_hot_action, reward, next_state, done,0])
        if done:
            if winner==1:
                for i in range(new_game_idx,len(self.replay_buffer),2):
                    self.replay_buffer[i][5]=1
                for i in range(new_game_idx+1,len(self.replay_buffer),2):
                    self.replay_buffer[i][5]=-1
            if winner==2:
                for i in range(new_game_idx+1,len(self.replay_buffer),2):
                    self.replay_buffer[i][5]=1
                for i in range(new_game_idx,len(self.replay_buffer),2):
                    self.replay_buffer[i][5]=-1




        if len(self.replay_buffer) > REPLAY_SIZE:
            self.replay_buffer.popleft()

        if len(self.replay_buffer) >BATCH_SIZE and done:
            self.train_Q_network()

    def train_Q_network(self):
        self.time_step += 1
        # Step 1: obtain random minibatch from replay memory
        minibatch = random.sample(self.replay_buffer, BATCH_SIZE)
        state_batch = [data[0] for data in minibatch]
        action_batch = [data[1] for data in minibatch]
        reward_batch = [data[2] for data in minibatch]
        next_state_batch = [data[3] for data in minibatch]
        winner_batch = [data[5] for data in minibatch]

        # Step 2: calculate y
        y_batch = []
        Q_value_batch = self.Q_value.eval(feed_dict={self.input_states: next_state_batch})
        #print('Q_value_batch:',Q_value_batch)
        for i in range(0, BATCH_SIZE):
            done = minibatch[i][4]
            if done:
                y_batch.append(reward_batch[i])
            else:
                y_batch.append(reward_batch[i] + GAMMA * np.max(Q_value_batch[i]))
                #print('np.max(Q_value_batch[i]):',np.max(Q_value_batch[i]))

        winner_batch = np.reshape(winner_batch, (-1, 1))
        print('winner_batch:',winner_batch)
        self.optimizer.run(feed_dict={
            self.y_input: y_batch,
            self.action_input: action_batch,
            self.input_states: state_batch,
            self.labels:winner_batch
        })
        if self.time_step%100==0:
            self.save_model('./current_policy.model')
            print('cost:',self.cost.eval(feed_dict={
            self.y_input: y_batch,
            self.action_input: action_batch,
            self.input_states: state_batch
        }))

    def DQN_ACTION(self,state,availbles):
        Q_value = self.Q_value.eval(feed_dict={
            self.input_states: [state]
        })[0]
        Q_value_index_list = np.argsort(Q_value)
        for i in range(1, len(Q_value_index_list) + 1):
            if Q_value_index_list[len(Q_value_index_list) - i] in availbles:
                return Q_value_index_list[len(Q_value_index_list) - i]

    def egreedy_action(self, state,availbles):
        #print('state:',state)

        #print('Q_value:',Q_value)

        if random.random() <= self.epsilon:
            randomvalue = random.sample(availbles, 1)[0]
            #print('randomvalue:',randomvalue)
            return randomvalue

        else:
            Q_value = self.Q_value.eval(feed_dict={
                self.input_states: [state]
            })[0]
            Q_value_index_list = np.argsort(Q_value)
            for i in range(1,len(Q_value_index_list)+1):
                if Q_value_index_list[len(Q_value_index_list)-i] in availbles:
                    return Q_value_index_list[len(Q_value_index_list)-i]




        self.epsilon -= (INITIAL_EPSILON - FINAL_EPSILON) /500 #5000

    def action(self, state):
        return np.argmax(self.Q_value.eval(feed_dict={
            self.state_input: [state]
        })[0])

    def weight_variable(self, shape):
        initial = tf.truncated_normal(shape)
        return tf.Variable(initial)

    def bias_variable(self, shape):
        initial = tf.constant(0.01, shape=shape)
        return tf.Variable(initial)

    def save_model(self, model_path):
        self.saver.save(self.session, model_path)

    def restore_model(self, model_path):
        self.saver.restore(self.session, model_path)

    def conv2d(self,x, W):
        return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

"""class DQN():
# DQN Agent
def __init__(self, env): #初始化

def create_Q_network(self): #创建Q网络

def create_training_method(self): #创建训练方法

def perceive(self,state,action,reward,next_state,done): #感知存储信息

def train_Q_network(self): #训练网络

def egreedy_action(self,state): #输出带随机的动作

def action(self,state): #输出动作"""